#include "scream.h"
#include "shared_mem.h"

CSharedMemory::CSharedMemory(IN PDEVICE_OBJECT DeviceObject) :
	m_bInitialized(FALSE),
	m_pIVShMemDeviceObject(NULL),
	m_pIVShMemFileObject(NULL),
	m_pDeviceObject(DeviceObject)
{
}

CSharedMemory::~CSharedMemory()
{
    ReleaseMMAP();
    if (m_pIVShMemFileObject) {
        ObDereferenceObject(m_pIVShMemFileObject);
        m_pIVShMemDeviceObject = NULL;
        m_pIVShMemFileObject = NULL;
    }
}

BOOLEAN CSharedMemory::Initialize()
{
    if (m_bInitialized) {
        return TRUE;
    }

    if (!m_pDeviceObject) {
        return FALSE;
    }

    // Enumerate IVSHMEM devices
    GUID interfaceClassGuid = {0xdf576976, 0x569d, 0x4672, 0x95, 0xa0, 0xf5, 0x7e, 0x4e, 0xa0, 0xb2, 0x10};
    PWSTR symbolicLinkList;
    NTSTATUS ntStatus;

    ntStatus = IoGetDeviceInterfaces(
        &interfaceClassGuid,
        NULL,
        0,
        &symbolicLinkList
    );
    if (!NT_SUCCESS(ntStatus) || NULL == symbolicLinkList) {
        return FALSE;
    }

    PFILE_OBJECT fileObj;
    PDEVICE_OBJECT devObj;
    UNICODE_STRING objName;

    KEVENT event;
    IO_STATUS_BLOCK ioStatus = { 0 };

    PIRP irp;
    UINT64 ivshmemSize = 0;

    for (PWSTR symbolicLink = symbolicLinkList;
        symbolicLink[0] != NULL && symbolicLink[1] != NULL;
        symbolicLink += wcslen(symbolicLink) + 1) {
        RtlInitUnicodeString(&objName, symbolicLink);

        ntStatus = IoGetDeviceObjectPointer(
                &objName,
                FILE_ALL_ACCESS,
                &fileObj,
                &devObj
            );
        if (!NT_SUCCESS(ntStatus)) {
            continue;
        }

        ivshmemSize = 0;
        KeInitializeEvent(&event, NotificationEvent, FALSE);
        irp = IoBuildDeviceIoControlRequest(
            IOCTL_IVSHMEM_REQUEST_SIZE,
            devObj,
            NULL, 0,
            &ivshmemSize, sizeof(ivshmemSize),
            FALSE,
            &event,
            &ioStatus
        );
        if (irp) {
            ntStatus = IoCallDriver(devObj, irp);
            if (ntStatus == STATUS_PENDING) {
                KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
            }
            //we suppose that we are the only 2*1MiB IVSHMEM device around
            if (ioStatus.Status == 0 && ivshmemSize == 2 * M_BYTES) {
                DPF_ENTER(("IVSHMEM Device Detected!\n"));

                m_pIVShMemDeviceObject = devObj;
                m_pIVShMemFileObject = fileObj;
                // we have found a suitable device, we skip the others, and we DON'T dereference devObj, it will be dereferenced in the destructor
                break;
            }
        }
        else {
            DPF_ENTER(("IVSHMEM: Failed to get the IRP for REQUEST_SIZE\n"));
        }

        ObDereferenceObject(fileObj);
    }

    ExFreePool(symbolicLinkList);

    return RequestMMAP();
}

PVOID CSharedMemory::GetSharedMemPtr()
{
    return m_bInitialized ? m_mmap.ptr : NULL;
}

UINT64 CSharedMemory::GetSharedMemSize()
{
    return m_bInitialized ? m_mmap.size : 0;
}

PDEVICE_OBJECT CSharedMemory::GetSharedDeviceObject()
{
    return m_bInitialized ? m_pIVShMemDeviceObject : NULL;
}

BOOLEAN CSharedMemory::RegisterIvshmemEventHandlerDispatch(PVOID Context, PIVSHMEM_EVENT_CALLBACK_FUNC func)
{
    if (!m_pIVShMemDeviceObject) {
        return FALSE;
    }

    KEVENT event;
    IVSHMEM_EVENT_CALLBACK evt_callback;
    IO_STATUS_BLOCK ioStatus = { 0 };
    evt_callback.pFunc = func;
    evt_callback.pContext = Context;
    KeInitializeEvent(&event, NotificationEvent, FALSE);
    PIRP irp = IoBuildDeviceIoControlRequest(
        IOCTL_IVSHMEM_REGISTER_EVENT_CB,
        m_pIVShMemDeviceObject,
        &evt_callback, sizeof(IVSHMEM_EVENT_CALLBACK),
        NULL, 0,
        FALSE,
        &event,
        &ioStatus
    );
    if (irp) {
        if (IoCallDriver(m_pIVShMemDeviceObject, irp) == STATUS_PENDING) {
            KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
            if (ioStatus.Status == 0) {
                return FALSE;
            }
        }
        else {
            DPF(D_ERROR, ("IRP for SET VOLUME failed with status %u\n", ioStatus.Status));
        }
    }
    return TRUE;
}

UINT32 CSharedMemory::GetVolumeFromBackend()
{
    if (!m_pIVShMemDeviceObject) {
        return FALSE;
    }

    IVSHMEM_REG inOut = { 0 };
    inOut.offset = IVSHMEM_BAR0_VOLUME_OFFSET;
    inOut.regSize = dword;

    if (!ReadWriteRegister(
        IOCTL_IVSHMEM_READ_REG,
        &inOut, sizeof(IVSHMEM_REG),
        &inOut, sizeof(IVSHMEM_REG))) {
        DPF(D_ERROR, ("IVSHMEM: Failed to call REQUEST_VOLUME\n"));
    }

    return inOut.dword;
}

VOID CSharedMemory::SetVolumeToBackend(UINT32 volume)
{
    if (!m_pIVShMemDeviceObject) {
        return;
    }

    IVSHMEM_REG in = { 0 };

    in.offset = IVSHMEM_BAR0_VOLUME_OFFSET;
    in.regSize = dword;
    in.dword = volume;
    if (!ReadWriteRegister(
        IOCTL_IVSHMEM_WRITE_REG,
        &in, sizeof(IVSHMEM_REG),
        NULL, 0)) {
        DPF(D_ERROR, ("IVSHMEM: Failed to call SET VOLUME\n"));
    }
}

VOID CSharedMemory::NotifyStart(BOOLEAN play, BOOLEAN start)
{
    IVSHMEM_REG in = { 0 };
    in.offset = IVSHMEM_BAR0_STATUS_OFFSET;
    in.regSize = dword;
    in.byte = play ? IVSHMEM_STATUS_PLAY_BIT : 0;
    in.byte |= start ? IVSHMEM_STATUS_START_BIT : 0;
    if (!ReadWriteRegister(
        IOCTL_IVSHMEM_WRITE_REG,
        &in, sizeof(IVSHMEM_REG),
        NULL, 0)) {
        DPF_ENTER(("IVSHMEM: Failed to call NotifyStart %s\n", play ? "play" : "capt"));
    }
}

BOOLEAN CSharedMemory::ReadWriteRegister(ULONG ioctl_code, PVOID in, ULONG insize, PVOID out, ULONG outsize)
{
    KEVENT event;
    IO_STATUS_BLOCK ioStatus = { 0 };
    KeInitializeEvent(&event, NotificationEvent, FALSE);
    if (m_pIVShMemDeviceObject != NULL) {
        PIRP irp = IoBuildDeviceIoControlRequest(
            ioctl_code,
            m_pIVShMemDeviceObject,
            in, insize,
            out, outsize,
            FALSE,
            &event,
            &ioStatus
        );
        if (irp) {
            if (IoCallDriver(m_pIVShMemDeviceObject, irp) == STATUS_PENDING) {
                KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
            }
            if (ioStatus.Status == 0) {
                return TRUE;
            }
            else {
                DPF_ENTER(("IVSHMEM: IRP for ReadWriteRegister failed with status %u\n", ioStatus.Status));
            }
        }
        else {
            DPF_ENTER(("IVSHMEM: Failed to get the IRP for ReadWriteRegister\n"));
        }
    }
    return FALSE;
}

UINT32 CSharedMemory::GetStatusFromBackend()
{
    if (!m_pIVShMemDeviceObject) {
        DPF(D_ERROR, ("GetStatusFromBackend: Failed for NULL m_pIVShMemDeviceObject\n"));
        return INVALID_IVSHMEM_STATUS;
    }

    IVSHMEM_REG inOut = { 0 };
    UINT32 status = INVALID_IVSHMEM_STATUS;
    inOut.offset = IVSHMEM_BAR0_STATUS_OFFSET;
    inOut.regSize = dword;

    if (ReadWriteRegister(IOCTL_IVSHMEM_READ_REG, &inOut, sizeof(IVSHMEM_REG), &inOut, sizeof(IVSHMEM_REG)) == TRUE) {
        status = inOut.dword;
    }

    return status;
}

BOOLEAN CSharedMemory::RequestMMAP()
{
    if (!m_pIVShMemDeviceObject) {
        return FALSE;
    }

    IO_STATUS_BLOCK ioStatus = { 0 };
    KEVENT event;
    KeInitializeEvent(&event, NotificationEvent, FALSE);

    IVSHMEM_MMAP_CONFIG mmapConfig;
    mmapConfig.cacheMode = IVSHMEM_CACHE_CACHED;

    PIRP irp = IoBuildDeviceIoControlRequest(
        IOCTL_IVSHMEM_REQUEST_KMAP,
        m_pIVShMemDeviceObject,
        &mmapConfig, sizeof(mmapConfig),
        &m_mmap, sizeof(IVSHMEM_MMAP),
        FALSE,
        &event,
        &ioStatus
    );
    if (irp) {
        if (IoCallDriver(m_pIVShMemDeviceObject, irp) == STATUS_PENDING) {
            KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        }
        if (ioStatus.Status == 0) {
            RtlZeroMemory(m_mmap.ptr, sizeof(IVSHMEM_SCREAM_HEADER));
            m_bInitialized = TRUE;
            return TRUE;
        }
        else {
            DPF_ENTER(("IVSHMEM: IRP for REQUEST_MMAP failed with status %u\n", ioStatus.Status));
        }
    }
    else {
        DPF_ENTER(("IVSHMEM: Failed to get the IRP for REQUEST_MMAP\n"));
    }
    return FALSE;
}

void CSharedMemory::ReleaseMMAP()
{
    if (!m_bInitialized) {
        return;
    }

    RtlZeroMemory(m_mmap.ptr, sizeof(IVSHMEM_SCREAM_HEADER));

    if (!m_pIVShMemDeviceObject) {
        return;
    }

    IO_STATUS_BLOCK ioStatus = { 0 };
    KEVENT event;
    KeInitializeEvent(&event, NotificationEvent, FALSE);
    PIRP irp = IoBuildDeviceIoControlRequest(
        IOCTL_IVSHMEM_RELEASE_MMAP,
        m_pIVShMemDeviceObject,
        NULL, 0,
        NULL, 0,
        FALSE,
        &event,
        &ioStatus
    );
    if (irp) {
        if (IoCallDriver(m_pIVShMemDeviceObject, irp) == STATUS_PENDING) {
            KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        }
        if (ioStatus.Status != 0) {
            DPF_ENTER(("IVSHMEM: IRP for RELEASE_MMAP failed with status %u\n", ioStatus.Status));
        }
    }
    else {
        DPF_ENTER(("IVSHMEM: Failed to get the IRP for RELEASE_MMAP\n"));
    }
    m_bInitialized = FALSE;
}
